# -*- coding: utf-8 -*-
# @author: HRUN

import psutil
import time
import threading
import logging
from datetime import datetime
from typing import Dict, List, Optional
import json

logger = logging.getLogger(__name__)


class SystemResourceMonitor:
    """系统资源监控器"""
    
    def __init__(self, report_id: Optional[int] = None, task_id: Optional[int] = None):
        self.report_id = report_id
        self.task_id = task_id
        self.monitoring = False
        self.monitor_thread = None
        self.interval = 5  # 监控间隔（秒）
        self.resource_data = []
        self.max_data_points = 100  # 最大保存数据点数
        
    def start_monitoring(self):
        """开始监控"""
        if self.monitoring:
            logger.warning("系统监控已在运行")
            return
            
        self.monitoring = True
        self.monitor_thread = threading.Thread(target=self._monitor_loop, daemon=True)
        self.monitor_thread.start()
        logger.info(f"系统资源监控已启动 - 报告ID: {self.report_id}")
        
    def stop_monitoring(self):
        """停止监控"""
        self.monitoring = False
        if self.monitor_thread:
            self.monitor_thread.join(timeout=10)
        logger.info(f"系统资源监控已停止 - 报告ID: {self.report_id}")
        
    def _monitor_loop(self):
        """监控循环"""
        while self.monitoring:
            try:
                # 收集资源数据
                resource_info = self.collect_system_info()
                self.resource_data.append(resource_info)
                
                # 限制数据点数量
                if len(self.resource_data) > self.max_data_points:
                    self.resource_data.pop(0)
                
                # 更新报告中的系统资源信息
                self.update_report_resources(resource_info)
                
                # 通过WebSocket广播系统资源信息
                self.broadcast_resource_update(resource_info)
                
                time.sleep(self.interval)
                
            except Exception as e:
                logger.error(f"系统监控循环出错: {e}")
                break
                
    def collect_system_info(self) -> Dict:
        """收集系统信息"""
        try:
            # CPU信息
            cpu_percent = psutil.cpu_percent(interval=1)
            cpu_count = psutil.cpu_count()
            cpu_freq = psutil.cpu_freq()
            
            # 内存信息
            memory = psutil.virtual_memory()
            swap = psutil.swap_memory()
            
            # 磁盘信息
            disk_usage = psutil.disk_usage('/')
            
            # 网络信息
            network_io = psutil.net_io_counters()
            
            # 进程信息
            process_count = len(psutil.pids())
            
            # 负载信息（Linux/Unix）
            load_avg = None
            try:
                load_avg = psutil.getloadavg()
            except AttributeError:
                # Windows 系统没有 getloadavg
                pass
            
            system_info = {
                'timestamp': datetime.now().isoformat(),
                'cpu': {
                    'percent': cpu_percent,
                    'count': cpu_count,
                    'frequency': {
                        'current': cpu_freq.current if cpu_freq else None,
                        'min': cpu_freq.min if cpu_freq else None,
                        'max': cpu_freq.max if cpu_freq else None
                    } if cpu_freq else None
                },
                'memory': {
                    'total': memory.total,
                    'available': memory.available,
                    'percent': memory.percent,
                    'used': memory.used,
                    'free': memory.free,
                    'buffers': getattr(memory, 'buffers', 0),
                    'cached': getattr(memory, 'cached', 0),
                    'shared': getattr(memory, 'shared', 0)
                },
                'swap': {
                    'total': swap.total,
                    'used': swap.used,
                    'free': swap.free,
                    'percent': swap.percent
                },
                'disk': {
                    'total': disk_usage.total,
                    'used': disk_usage.used,
                    'free': disk_usage.free,
                    'percent': (disk_usage.used / disk_usage.total) * 100 if disk_usage.total > 0 else 0
                },
                'network': {
                    'bytes_sent': network_io.bytes_sent,
                    'bytes_recv': network_io.bytes_recv,
                    'packets_sent': network_io.packets_sent,
                    'packets_recv': network_io.packets_recv,
                    'errin': network_io.errin,
                    'errout': network_io.errout,
                    'dropin': network_io.dropin,
                    'dropout': network_io.dropout
                },
                'processes': {
                    'count': process_count
                },
                'load_average': load_avg
            }
            
            return system_info
            
        except Exception as e:
            logger.error(f"收集系统信息失败: {e}")
            return {
                'timestamp': datetime.now().isoformat(),
                'error': str(e)
            }
    
    def update_report_resources(self, resource_info: Dict):
        """更新报告中的资源信息"""
        if not self.report_id:
            return
            
        try:
            import os
            import django
            # 统一通过ENV变量判断
            if "DJANGO_SETTINGS_MODULE" not in os.environ:
                env = os.environ.get("ENV")
                if env == "production":
                    settings_module = "primaryApp.settings.pro"
                else:
                    settings_module = "primaryApp.settings.dev"
                os.environ.setdefault("DJANGO_SETTINGS_MODULE", settings_module)
            django.setup()
            
            from performance.models import TaskReport
            import json
            
            try:
                report = TaskReport.objects.get(id=self.report_id)
            except TaskReport.DoesNotExist:
                logger.warning(f"TaskReport id={self.report_id} 不存在，自动停止监控。")
                self.monitoring = False
                return
            
            # 只有在报告状态为执行中时才更新CPU和内存值
            if report.reportStatus == '1':
                # 更新当前CPU和内存使用率（使用最新值）
                if 'cpu' in resource_info and 'memory' in resource_info:
                    current_cpu = resource_info['cpu']['percent']
                    current_memory = resource_info['memory']['percent']
                    
                    # 直接使用当前最新值
                    report.avgCpu = current_cpu
                    report.avgMemory = current_memory
            
            # 无论报告状态如何，都将系统资源数据添加到报告结果中（用于历史记录）
            try:
                # 获取当前报告结果
                report_result = {}
                if report.reportResult:
                    try:
                        report_result = json.loads(report.reportResult)
                    except json.JSONDecodeError:
                        report_result = {}
                
                # 确保stats_history存在
                if 'stats_history' not in report_result:
                    report_result['stats_history'] = []
                
                # 为最新的统计数据添加系统信息
                if report_result['stats_history']:
                    latest_stat = report_result['stats_history'][-1]
                    latest_stat['system_info'] = resource_info
                else:
                    # 如果没有统计数据，创建一个只包含系统信息的记录
                    report_result['stats_history'].append({
                        'timestamp': resource_info.get('timestamp', ''),
                        'system_info': resource_info
                    })
                
                # 限制历史记录数量，避免数据过大
                if len(report_result['stats_history']) > 100:
                    report_result['stats_history'] = report_result['stats_history'][-100:]
                
                # 更新报告结果
                report.reportResult = json.dumps(report_result)
                
            except Exception as e:
                logger.error(f"更新报告系统资源历史数据失败: {e}")
            
            # 只有在报告状态为执行中或者刚刚更新了报告结果时才保存
            report.save()
                
        except Exception as e:
            logger.error(f"更新报告资源信息失败: {e}")
    
    def broadcast_resource_update(self, resource_info: Dict):
        """广播资源更新"""
        if not self.task_id:
            return
            
        try:
            import threading
            import asyncio
            from performance.consumers import RealtimeDataBroadcaster
            
            def run_async_broadcast():
                """在新线程中运行异步广播"""
                try:
                    # 检查是否有现有的事件循环
                    try:
                        loop = asyncio.get_running_loop()
                        # 如果有运行中的循环，使用 call_soon_threadsafe
                        future = asyncio.run_coroutine_threadsafe(
                            RealtimeDataBroadcaster.broadcast_performance_update(
                                self.task_id, 
                                {
                                    'type': 'system_resources',
                                    'data': resource_info
                                }
                            ),
                            loop
                        )
                        # 不等待结果，避免阻塞
                    except RuntimeError:
                        # 没有运行中的事件循环，创建新的
                        loop = asyncio.new_event_loop()
                        asyncio.set_event_loop(loop)
                        try:
                            loop.run_until_complete(
                                RealtimeDataBroadcaster.broadcast_performance_update(
                                    self.task_id, 
                                    {
                                        'type': 'system_resources',
                                        'data': resource_info
                                    }
                                )
                            )
                        finally:
                            loop.close()
                except Exception as e:
                    logger.debug(f"异步广播线程执行失败: {e}")
            
            # 在后台线程中执行异步操作
            thread = threading.Thread(target=run_async_broadcast, daemon=True)
            thread.start()
            
        except Exception as e:
            logger.debug(f"WebSocket广播资源信息失败: {e}")
    
    def get_average_resources(self) -> Dict:
        """获取平均资源使用情况"""
        if not self.resource_data:
            return {}
        
        try:
            cpu_values = [data['cpu']['percent'] for data in self.resource_data if 'cpu' in data and data['cpu']['percent'] is not None]
            memory_values = [data['memory']['percent'] for data in self.resource_data if 'memory' in data]
            disk_values = [data['disk']['percent'] for data in self.resource_data if 'disk' in data]
            
            return {
                'avg_cpu': sum(cpu_values) / len(cpu_values) if cpu_values else 0,
                'max_cpu': max(cpu_values) if cpu_values else 0,
                'min_cpu': min(cpu_values) if cpu_values else 0,
                'avg_memory': sum(memory_values) / len(memory_values) if memory_values else 0,
                'max_memory': max(memory_values) if memory_values else 0,
                'min_memory': min(memory_values) if memory_values else 0,
                'avg_disk': sum(disk_values) / len(disk_values) if disk_values else 0,
                'data_points': len(self.resource_data)
            }
            
        except Exception as e:
            logger.error(f"计算平均资源失败: {e}")
            return {}
    
    def get_resource_history(self, limit: int = 50) -> List[Dict]:
        """获取资源历史数据"""
        return self.resource_data[-limit:] if self.resource_data else []


class ProcessResourceMonitor:
    """进程资源监控器"""
    
    def __init__(self, process_name: str = "locust"):
        self.process_name = process_name
        self.processes = []
        
    def find_target_processes(self) -> List[psutil.Process]:
        """查找目标进程"""
        processes = []
        try:
            for proc in psutil.process_iter(['pid', 'name', 'cmdline']):
                try:
                    if self.process_name.lower() in proc.info['name'].lower():
                        processes.append(psutil.Process(proc.info['pid']))
                except (psutil.NoSuchProcess, psutil.AccessDenied):
                    continue
        except Exception as e:
            logger.error(f"查找进程失败: {e}")
        
        return processes
    
    def get_process_resources(self) -> Dict:
        """获取进程资源使用情况"""
        processes = self.find_target_processes()
        
        if not processes:
            return {'error': f'未找到 {self.process_name} 进程'}
        
        total_cpu = 0
        total_memory = 0
        total_threads = 0
        process_count = len(processes)
        
        process_details = []
        
        for proc in processes:
            try:
                cpu_percent = proc.cpu_percent()
                memory_info = proc.memory_info()
                memory_percent = proc.memory_percent()
                num_threads = proc.num_threads()
                
                total_cpu += cpu_percent
                total_memory += memory_percent
                total_threads += num_threads
                
                process_details.append({
                    'pid': proc.pid,
                    'cpu_percent': cpu_percent,
                    'memory_rss': memory_info.rss,
                    'memory_vms': memory_info.vms,
                    'memory_percent': memory_percent,
                    'num_threads': num_threads,
                    'status': proc.status()
                })
                
            except (psutil.NoSuchProcess, psutil.AccessDenied):
                continue
        
        return {
            'process_count': process_count,
            'total_cpu_percent': total_cpu,
            'avg_cpu_percent': total_cpu / process_count if process_count > 0 else 0,
            'total_memory_percent': total_memory,
            'avg_memory_percent': total_memory / process_count if process_count > 0 else 0,
            'total_threads': total_threads,
            'process_details': process_details,
            'timestamp': datetime.now().isoformat()
        }


# 全局监控实例字典，按报告ID索引
system_monitors = {}


def start_system_monitoring(report_id: int, task_id: int):
    """启动系统监控"""
    global system_monitors
    
    # 检查是否已有该报告的监控实例
    if report_id in system_monitors and system_monitors[report_id].monitoring:
        logger.warning(f"报告 {report_id} 的系统监控已在运行")
        return system_monitors[report_id]
        
    # 创建新的监控实例
    monitor = SystemResourceMonitor(report_id, task_id)
    monitor.start_monitoring()
    system_monitors[report_id] = monitor
    return monitor


def stop_system_monitoring(report_id: Optional[int] = None):
    """停止系统监控"""
    global system_monitors
    
    if report_id is not None:
        # 停止指定报告的监控
        if report_id in system_monitors:
            system_monitors[report_id].stop_monitoring()
            del system_monitors[report_id]
            logger.info(f"已停止报告 {report_id} 的系统监控")
    else:
        # 停止所有监控
        for monitor_id, monitor in list(system_monitors.items()):
            monitor.stop_monitoring()
            logger.info(f"已停止报告 {monitor_id} 的系统监控")
        system_monitors.clear()


def get_current_system_monitor(report_id: Optional[int] = None):
    """获取当前系统监控器"""
    if report_id is not None:
        return system_monitors.get(report_id)
    
    # 如果没有指定报告ID，返回第一个监控器
    if system_monitors:
        return next(iter(system_monitors.values()))
    
    return None